#include "ast.hpp"
#include "lexer.hpp"
#include "visitor.hpp"
#include "type.hpp"
#include "symtable.hpp"

#include <algorithm>
#include <cassert>
#include <iostream>
using namespace std;

/* constructors. */
ListNode::ListNode() {
    _node_list = new vector<Node*>();
}

void ListNode::add(Node* n) {
    _node_list->push_back(n);
    n->set_parent(this);
}

void ListNode::remove(Node* n) {
    auto it = find(_node_list->begin(), _node_list->end(), n);

    if (it != _node_list->end()) {
        _node_list->erase(it);
    }
}

void ListNode::insert(Node* n, Node* t) {
    auto it = find(_node_list->begin(), _node_list->end(), n);

    if (it != _node_list->end()) {
        _node_list->insert(it, t);
    }
    else {
        _node_list->push_back(t);
    }
}

StmtsNode::StmtsNode() {
}

Node* TupleNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* TypeTupleNode::accept(Visitor* v) {
    return v->visit(this);
}

VarNode::VarNode(Token* t) {
    _name = string(t->_value, t->_length);
}

FuncDefNode::FuncDefNode(Token* t) {
    _name = string(t->_value, t->_length);
    _generic_args = nullptr;
    _sym_table = nullptr;
    _ret_type = nullptr;
    _return_type_node = nullptr;
    _args = nullptr;
    _body = nullptr;
    _poly_params = nullptr;
    _type = nullptr;
}

TypeTable* FuncDefNode::type_table() {
    return _sym_table->type_table();
}

void BinaryOp::set_left(Node* left) {
    _left = left;
    _left->set_parent(this);
}

void BinaryOp::set_right(Node* right) {
    _right = right;
    _right->set_parent(this);
}

/**
 * visitor pattern for AST.
 */
Node* AssignNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* VarNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* CmpNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* ConstInt::accept(Visitor* v) {
    return v->visit(this);
}

Node* ConstArray::accept(Visitor* v) {
    return v->visit(this);
}

ConstString::ConstString(int length, const char* s) {
    _value = new char[length + 1];
    strncpy(_value, s, length);
    _value[length] = '\0';
    _length = length;
}

Node* ConstDouble::accept(Visitor* v) {
    return v->visit(this);
}

ConstString::~ConstString() {
    delete[] _value;
    _value = nullptr;
}

Node* ConstString::accept(Visitor* v) {
    return v->visit(this);
}

Node* PrintNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* ReturnNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* ImportNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* IfNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* WhileNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* ParamDefNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* VarDefNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* ListNode::accept(Visitor* v) {
    return v->visit(this);
}


Node* StmtsNode::accept(Visitor *v) {
    return v->visit(this);
}

Node* BinaryOp::accept(Visitor* v) {
    return v->visit(this);
}

Node* LogicNotNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* FuncDefNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* RecordDefNode::accept(Visitor* v) {
    return v->visit(this);
}

void RecordDefNode::add_method(FuncDefNode* method) {
    if (!_methods) {
        _methods = new ListNode();
    }
    _methods->add(method);
}

Node* ImplementNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* TraitDefNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* CallNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* SubscrNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* DotNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* SequenceNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* LambdaDef::accept(Visitor* visitor) {
    return visitor->visit(this);
}

Node* TypeFuncDefNode::accept(Visitor* visitor) {
    return visitor->visit(this);
}

Node* TypeVarDefNode::accept(Visitor* visitor) {
    return visitor->visit(this);
}

Node* ArrayTypeNode::accept(Visitor* visitor) {
    return visitor->visit(this);
}

Node* ArrowTypeNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* TypeConstInt::accept(Visitor* v) {
    return v->visit(this);
}

Node* TypeCallNode::accept(Visitor* v) {
    return v->visit(this);
}

Node* TypeVarNode::accept(Visitor* v) {
    return v->visit(this);
}

Type* Node::get_type(TypeTable* type_table) {
    if (!_type) {
        return nullptr;
    }

    if (_type->is_type_var()) {
        TypeVar* tv = dynamic_cast<TypeVar*>(_type);
        Type* sbt = type_table->get_substitution(tv->name());
        if (sbt) {
            // do not assign sbt to _type because we do not remember
            // the substitution type of type var.
            return sbt;
        }
    }

    return _type;
}

Type* TypeTupleNode::get_type(TypeTable* type_table) {
    vector<Type*> types;
    for (auto it = _args->node_list()->begin();
        it != _args->node_list()->end(); it++) {
        types.push_back((*it)->get_type(type_table));
    }
    return new TupleType(types);
}

Type* TypeVarNode::get_type(TypeTable* type_table) {
    Type* t = type_table->get_substitution(name());
    if (t) {
        return t;
    }

    t = type_table->get_expanded_type(name());
    if (t) {
        return t;
    }

    t = new TypeVar(name(), type_table);
    type_table->add(name(), t);

    return t;
}

Type* ArrowTypeNode::get_type(TypeTable* type_table) {
    if (_type) {
        return _type;
    }

    Type* src_type = _src->get_type(type_table);
    Type* dst_type = _dst->get_type(type_table);
    _type = new ArrowType(src_type, dst_type);
    
    return _type;
}

Type* ConstInt::get_type(TypeTable* type_table) {
    return IntType::get_instance();
}

Type* ConstString::get_type(TypeTable* type_table) {
    return StringType::get_instance();
}

Type* ConstDouble::get_type(TypeTable* type_table) {
    return DoubleType::get_instance();
}

Type* TypeCallNode::get_type(TypeTable* type_table) {
    if (_type) {
        return _type;
    }

    Type* callee_type = _name->get_type(type_table);
    assert(callee_type->is_type_func());
    TypeFunction* tf = dynamic_cast<TypeFunction*>(callee_type);
    if (tf->args().size() != _args->size()) {
        throw "Type mismatch";
    }

    for (int i = 0; i < _args->size(); i++) {
        Type* arg_type = _args->get(i)->get_type(type_table);
        TypeVar* tv = dynamic_cast<TypeVar*>(tf->args()[i]);
        type_table->add_substitution(tv->name(), arg_type);
    }

    _type = tf->expand(type_table);
    return _type;
}

Node* BuiltinConst::accept(Visitor* visitor) {
    return visitor->visit(this);
}

/*
 * some initialization.
 */

BuiltinConst* BuiltinConst::UnitNode = new BuiltinConst(BuiltinConst::UNIT);
BuiltinConst* BuiltinConst::TrueNode = new BuiltinConst(BuiltinConst::TRUE);
BuiltinConst* BuiltinConst::FalseNode = new BuiltinConst(BuiltinConst::FALSE);

